package skewreduce.lib;

import java.util.ArrayList;
import java.util.IdentityHashMap;
import java.util.Iterator;
import java.util.List;

public final class PartitionInfo implements Iterable<Partition> {
    final Partition root;
    List<Partition> leafNodes;
    IdentityHashMap<Partition,Integer> partitionMap = new IdentityHashMap<Partition,Integer>();

    public PartitionInfo(Partition bound) {
        this.root = bound;
        build();
    }

    private void traverse(Partition n) {
        if ( n.isLeaf() ) {
            partitionMap.put(n,leafNodes.size());
            leafNodes.add(n);
        } else {
            if ( n.hasLeft() ) {
                traverse(n.getLeft());
            }
            if ( n.hasRight() ) {
                traverse(n.getRight());
            }
        }
    }

    private void build() {
        leafNodes = new ArrayList<Partition>();
        traverse(root);
    }

    public Partition get(int i) {
    	return leafNodes.get(i);
    }
    
    public int getPartition(Partition c) {
        return partitionMap.get(c);
    }

    public <T extends ITuple> int getPartition(T p) {
        return partitionMap.get(root.getLeaf(p));
    }

    public int size() { return leafNodes.size(); }

    public String toString() {
        StringBuilder buf = new StringBuilder(1024);
        int i = 0;
        for ( Partition c : leafNodes ) {
            buf.append("partition ");
            buf.append(i++);
            buf.append(':');
            buf.append(c);
            buf.append('\n');
        }
        return buf.toString();
    }

	@Override
	public Iterator<Partition> iterator() {
		return leafNodes.iterator();
	}
}
